package com.geor.gcv.ict.classifier;

import com.geor.common.log.Loggers;
import com.geor.gcv.core.mat.Mats;
import com.geor.gcv.ict.data.ClassifierResult;
import org.apache.commons.io.IOUtils;
import org.nutz.json.Json;
import org.nutz.lang.Files;
import org.nutz.lang.Lang;
import org.nutz.lang.encrypt.DESPlus;
import org.opencv.core.Mat;
import org.opencv.core.Size;
import org.slf4j.Logger;
import org.tensorflow.Graph;
import org.tensorflow.Session;
import org.tensorflow.Session.Runner;
import org.tensorflow.Tensor;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileReader;
import java.util.*;

/**
 * java 调用tensorflow 接口，识别图片
 *
 * @author geor.lusongnan
 */
public class ClassifierHandler {
    static Logger log = Loggers.get();
    private Map<String, String> dictory = null;
    private Graph graph = new Graph();
    private Session session;

    public ClassifierHandler() {
    }

    /**
     * 分类器加载需要的字典文件
     *
     * @param dict
     * @throws Exception
     */
    public void loadDict(String dict) throws Exception {
        try {
            dictory = Json.fromJsonAsMap(String.class, new FileReader(new File(dict)));
            log.debug(String.format("[%s] load success %d", dict, dictory.size()));
        } catch (Exception e) {
            String msg = DESPlus.decrypt(Files.read(new File(dict)));
            try {
                dictory = Json.fromJsonAsMap(String.class, msg);
                log.debug(String.format("[%s] load success %d", dict, dictory.size()));
            } catch (Exception e1) {
                throw new Exception(String.format("[ %s ]字典文件读取失败", dict));
            }
        }
    }

    public void loadModel(String path) throws Exception {
        byte[] graphBytes = null;
        try {
            graphBytes = IOUtils.toByteArray(new FileInputStream(path));
            graph.importGraphDef(graphBytes);
            session = new Session(graph);
            log.debug(String.format("[%s] load success", path));
        } catch (Exception e) {
            log.error(String.format("load [%s] fail", path), e);
        }
    }

    public List<ClassifierResult> predicts(List<Mat> mats, int width, int height, int size) {
        List<List<PredictWapper>> pws = _predicts(mats, width, height, size);
        if (Lang.isEmpty(pws)) {
            return null;
        }
        List<ClassifierResult> ccsResults = new ArrayList<ClassifierResult>();
        for (List<PredictWapper> results : pws) {
            ClassifierResult ccsResult = new ClassifierResult();

            String[] responses = new String[results.size()];
            float[] similarites = new float[results.size()];

            for (int i = 0; i < results.size(); i++) {
                PredictWapper result = results.get(i);
                responses[i] = result.val;
                similarites[i] = result.predictVal;
            }

            ccsResult.responses = responses;
            ccsResult.similarites = similarites;

            ccsResults.add(ccsResult);
        }
        return ccsResults;
    }

    public synchronized List<List<PredictWapper>> _predicts(List<Mat> mats, int width, int height, int size) {
        if (Lang.isEmpty(dictory)) {
            log.warn("字典为空, 尝试重新加载字典文件");
        }

        if (session == null) {
            throw new RuntimeException("未加载模型或加载模型失败");
        }

        for (int i = 0; i < mats.size(); i++) {
            Mat mat = mats.get(i);
            Mat resize = Mats.resize(mat, new Size(width, height));
            mats.set(i, resize);
        }

        Tensor<Float> image = null;
        Tensor<?> keepProb = Tensor.create(1.0f);
        List<Tensor<?>> tf = null;
        try {
            image = TensorFlows.doToTensorDatas(mats);
            Runner runner = session.runner();
            runner.feed("images_bacth", image);
            runner = runner.feed("keep_prob", keepProb);
            runner = runner.fetch("output");
            tf = runner.run();
            Tensor<?> tensor = tf.get(0);
            long[] rshape = tensor.shape();
            float[][] f = tensor.copyTo(new float[mats.size()][(int) rshape[1]]);

            List<List<PredictWapper>> pws = new ArrayList<List<PredictWapper>>();
            for (int i = 0; i < f.length; i++) {
                List<PredictWapper> results = new ArrayList<PredictWapper>();
                for (int j = 0; j < f[i].length; j++) {
                    if (f[i][j] * 10000 > 0) {
                        results.add(new PredictWapper(f[i][j], j));
                    }
                }
                Collections.sort(results, new ComparatorImpl());
                pws.add(results);
            }

            for (int i = 0; i < pws.size(); i++) {
                List<PredictWapper> tmps = new ArrayList<PredictWapper>();
                List<PredictWapper> results = pws.get(i);
                for (int j = 0; j < results.size() && j < size; j++) {
                    PredictWapper predictWapper = results.get(j);
                    if (Lang.isEmpty(dictory)) {
                        predictWapper.val = String.valueOf(predictWapper.predictIndex);
                    } else {
                        predictWapper.val = dictory.get(String.valueOf(predictWapper.predictIndex));
                    }
                    tmps.add(predictWapper);
                }
                pws.set(i, tmps);
            }
            return pws;
        } finally {
            if (image != null) {
                image.close();
            }
            if (keepProb != null) {
                keepProb.close();
            }
            if (tf != null) {
                for (Tensor<?> t : tf) {
                    t.close();
                }
            }
        }
    }

    public synchronized List<PredictWapper> predict(Mat mat, int width, int height) {
        return predict(mat, width, height, 15);
    }

    public synchronized List<PredictWapper> predict(Mat mat, int width, int height, int size) {
        List<List<PredictWapper>> pws = _predicts(Arrays.asList(mat), width, height, size);
        if (Lang.isEmpty(pws)) {
            return null;
        }
        return pws.get(0);
    }

    private class ComparatorImpl implements Comparator<PredictWapper> {

        @Override
        public int compare(PredictWapper o1, PredictWapper o2) {
            if (o1.predictVal > o2.predictVal) {
                return -1;
            }

            if (o1.predictVal < o2.predictVal) {
                return 1;
            }

            return 0;
        }

    }

    public void relese() {
        if (session != null) {
            session.close();
            session = null;
        }
        dictory = null;
    }

}
